import numpy as np
from scipy.spatial.transform import Rotation as R
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def generate_platform_points(diameter, num_points=6, offset_angle=0):
    """生成闭合的平台铰链点坐标"""
    radius = diameter / 2
    angles = np.linspace(0, 2*np.pi, num_points, endpoint=False) + offset_angle
    points = np.array([(radius*np.cos(theta), radius*np.sin(theta), 0) for theta in angles])
    return np.vstack([points, points[0]])  # 添加闭合点

def calculate_kinematics(lower_points, upper_local, position, direction, twist_angle):
    """计算支腿长度和位姿矩阵"""
    direction = direction / np.linalg.norm(direction)
    align_rot, _ = R.align_vectors([[0,0,1]], [direction])
    twist_rot = R.from_rotvec(twist_angle * direction)
    rotation = (twist_rot * align_rot).as_matrix()
    upper_global = (rotation @ upper_local[:-1].T).T + position  # 排除闭合点
    return [np.linalg.norm(upper_global[i]-lower_points[i]) for i in range(6)], upper_global, rotation

def plot_platform(lower_pts, upper_pts, lower_dia, upper_dia, position, rotation):
    """带正确圆形连接的可视化"""
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 下平台圆形
    theta = np.linspace(0, 2*np.pi, 100)
    lower_circle = lower_dia/2 * np.column_stack([np.cos(theta), np.sin(theta), np.zeros_like(theta)])
    ax.plot(lower_circle[:,0], lower_circle[:,1], lower_circle[:,2], 'r--', alpha=0.6, linewidth=1.5)
    
    # 上平台圆形（正确变换）
    upper_circle_local = generate_platform_points(upper_dia, 6, np.pi/6)[:-1]  # 匹配铰链点生成方式
    upper_circle_global = (rotation @ upper_circle_local.T).T + position
    upper_circle_global = np.vstack([upper_circle_global, upper_circle_global[0]])  # 闭合圆环
    ax.plot(upper_circle_global[:,0], upper_circle_global[:,1], upper_circle_global[:,2], 
            'b-', alpha=0.8, linewidth=1.5)
    
    # 绘制支腿和连接点
    ax.scatter(lower_pts[:,0], lower_pts[:,1], lower_pts[:,2], c='r', s=50, depthshade=False)
    ax.scatter(upper_pts[:,0], upper_pts[:,1], upper_pts[:,2], c='b', s=50, depthshade=False)
    for i in range(6):
        ax.plot([lower_pts[i,0], upper_pts[i,0]], 
                [lower_pts[i,1], upper_pts[i,1]], 
                [lower_pts[i,2], upper_pts[i,2]], 
                'g-', linewidth=2, alpha=0.7)
    
    # 设置坐标轴
    ax.set_xlabel('X (mm)'), ax.set_ylabel('Y (mm)'), ax.set_zlabel('Z (mm)')
    ax.xaxis.pane.set_alpha(0.5), ax.yaxis.pane.set_alpha(0.5), ax.zaxis.pane.set_alpha(0.5)
    plt.title('Stewart Platform Simulation')
    plt.show()

# 参数输入
lower_dia = float(input("下平台直径(mm): ")) 
upper_dia = float(input("上平台直径(mm): "))
min_len = float(input("支腿最小长度(mm): "))
max_len = float(input("支腿最大长度(mm): "))

# 生成平台结构（包含闭合点）
lower_pts = generate_platform_points(lower_dia)[:-1]  # 下平台实际连接点
upper_pts_local = generate_platform_points(upper_dia, offset_angle=np.pi/6)  # 上平台局部坐标

# 目标位姿输入
pos = np.array([float(input(f"输入{i}坐标(mm): ")) for i in ['X', 'Y', 'Z']])
direction = np.array([float(input(f"方向向量{i}分量: ")) for i in ['X', 'Y', 'Z']])
twist = np.radians(float(input("扭转角度(度): ")))

# 运动学计算
lengths, upper_pts, rot_mat = calculate_kinematics(lower_pts, upper_pts_local, pos, direction, twist)

# 结果验证
valid = all(min_len <= l <= max_len for l in lengths)
print(f"可达性: {'是' if valid else '否'}")
print(f"支腿长度(mm): {np.round(lengths, 2)}")

# 可视化
plot_platform(lower_pts, upper_pts, lower_dia, upper_dia, pos, rot_mat)
